package com.dyava.commons.token_bucket.local;

import com.dyava.commons.token_bucket.BToken;
import com.dyava.commons.token_bucket.TokenBucket;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;

import java.util.concurrent.TimeUnit;

public class LocalTokenBucket extends TokenBucket {
    Cache<String,Object> cache;


    public LocalTokenBucket(String name, int limitNum, int limitDen){
        super(name,limitNum,limitDen);
        cache = CacheBuilder.newBuilder().expireAfterWrite(limitDen*2, TimeUnit.SECONDS).build(); //2个周期都没人用，就删除
    }


    @Override
    public boolean limit(Object id) {
        String idStr = id==null? NULL : id.toString();
        BToken token = (BToken)cache.getIfPresent(idStr);
        if(token==null){
            token = new BToken(limitNum);
        }
        int expectNum = calculateExpectNum(token);
        if(expectNum > 0){
            expectNum--;
            token.setLastExecuteTime(System.currentTimeMillis());
            token.setNum(expectNum);
            cache.put(idStr,token);
            return true;
        }else{
            return false;
        }
    }

}
